import random
import numpy as np
import torch
from torch import nn


class ReplayBuffer:
    def __init__(self, memory_size):
        self.memory_size = memory_size
        self.memory = []
        self.position = 0

    def record(self, state, action, reward, next_state, done):
        if not self.is_ready():
            self.memory.append([state, action, reward, next_state, done])
            self.position = (self.position + 1) % self.memory_size
        else:
            self.memory[self.position] = [state, action, reward, next_state, done]
            self.position = (self.position + 1) % self.memory_size

    def is_ready(self):
        return len(self.memory) == self.memory_size

    def sample(self, batch_size):
        if not self.is_ready():
            raise ValueError("Not enough samples in memory")
        batch = random.sample(self.memory, batch_size)
        state, action, reward, next_state, done = map(
            lambda x: np.stack(x).reshape(batch_size, -1), zip(*batch)
        )
        return state, action, reward, next_state, done


class Network(nn.Module):
    """
    input: state->s
    output: qa(s,a1),qa(s,a2),...,qa(s,an)
    """

    def __init__(self, input: int, output: int) -> None:
        super(Network, self).__init__()
        self.fc1 = nn.Linear(input, 64)
        self.fc4 = nn.Linear(64, 16)
        self.fc5 = nn.Linear(16, output)

    def forward(self, state: torch.Tensor):
        state = nn.functional.relu(self.fc1(state))
        state = nn.functional.relu(self.fc4(state))
        state = self.fc5(state)
        return state

    def save(self, path: str):
        torch.save(self.state_dict(), path)

    def load(self, path: str):
        self.load_state_dict(torch.load(path))


class DQN:
    def __init__(
        self,
        state_shape_len: int,
        action_len: int,
        lr: float,
        gamma: float,
        epsilon: float,
        batch_size: int,
        memory_size: int,
        update_freq: int,
    ):
        self.state_shape_len = state_shape_len
        self.action_len = action_len
        self.lr = lr
        self.gamma = gamma
        self.epsilon = epsilon
        self.batch_size = batch_size
        self.update_freq = update_freq
        self.memory = ReplayBuffer(memory_size)
        self.QNet = Network(state_shape_len, action_len)
        self.QTargetNet = Network(state_shape_len, action_len)
        self.QTargetNet.load_state_dict(self.QNet.state_dict())
        self.QNet.train()
        self.QTargetNet.eval()
        self.criterion = torch.nn.MSELoss()
        self.optimizer = torch.optim.Adam(self.QNet.parameters(), lr=self.lr)
        self.freq_count = 0

    def sample_action(self, state):
        with torch.no_grad():
            if np.random.random() < self.epsilon:
                return np.random.randint(self.action_len)
            else:
                stateTensor = torch.tensor(state, dtype=torch.float32).reshape(
                    -1, self.state_shape_len
                )
                return self.QNet(stateTensor).max(dim=1)[1].item()

    def update(self, state, action, reward, next_state, done):
        if not done:
            self.memory.record(state, action, reward, next_state, done)
        elif self.memory.is_ready():
            self.learn()

    def learn(self):
        state, action, reward, next_state, done = self.memory.sample(self.batch_size)
        stateTensor = torch.tensor(state, dtype=torch.float32)
        next_stateTensor = torch.tensor(next_state, dtype=torch.float32)
        rewardTensor = torch.tensor(reward, dtype=torch.float32)
        actionTensor = torch.tensor(action, dtype=torch.long)
        doneTensor = torch.tensor(done, dtype=torch.float32)

        with torch.no_grad():
            # targetQa[[q(s1,a1),q(s1,a2)],[q(s2,a1),q(s2,a2)]]
            targetQ = self.QTargetNet(stateTensor)
            targetNextQa = self.QTargetNet(next_stateTensor).max(dim=1)[0]
            # 根据采样的状态和动作，更新targetQ，但是只有部分的Qa值需要更新
            # batchsize 个状态，最多更新 batchsize 个Qa值，其他不变
            targetQ[range(self.batch_size), actionTensor] = (
                rewardTensor + (1.0 - doneTensor) * targetNextQa * self.gamma
            )
        predictQ = self.QNet(stateTensor)

        loss = self.criterion(predictQ, targetQ)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        if self.freq_count % self.update_freq == 0:
            self.QTargetNet.load_state_dict(self.QNet.state_dict())
        self.freq_count += 1

    def printQsTable(self):
        with torch.no_grad():
            for i in range(16):
                qsa = self.QTargetNet(torch.tensor([[float(i)]]))
                for j in range(self.action_len):
                    print(f"{qsa[0][j].item():.2f}", end=" ")
                print()
